PBDS Set에 관한 간단한 정리. 알게 되었는데 정리 안하면 나중에 다시 맨땅에서 구글링 해야하므로 한번 알게 된 내용 정리하고 가보려고 한다.

• gnu official design manual: http://gcc.gnu.org/onlinedocs/libstdc++/manual/policy_data_structures_design.html
• 관련 Codeforces 글: https://codeforces.com/blog/entry/11080

PBDS Set?

PBDS는 Policy based data structures의 약자로, g++에서 ext/pb_ds 아래에 위치하는 데이터 구조를 칭한다. (MSVC에서는 사용이 어려울 듯) 보통 PS에서 이것을 사용하는 의의 중 하나는 기존 STL Set에서 지원하지 않는 K번째 크기가 큰(또는 작은) 원소를 빠르게 찾을 수 있게 하는 데이터 구조를 사용하기 위해 사용하는 것이다.

사용 방법

일반 set을 사용하는 느낌으로 사용할 수 있게 define해서 사용한다. 글 하단에 예시가 있다. 여기서는 ordered_set으로 사용한다.

일반 set에서 사용되는 method를 사용함과 더불어(심지어 iterator도 사용가능함을 확인했다) 추가로 사용되는 method는 find_by_order(k)와 order_of_key(key)이다. 이름 그대로, k번째 원소를 $O(\log n)$ 에 찾을 수 있게 하거나, key값을 주었을 때 그것이 몇번째 원소인지를 $O(\log n)$에 알게 해준다. 일반 set에서 같은 일을 할 수 있으려면 $O(n)$의 시간복잡도였던 것을 감안하면 큰 차이이다.

ordered_set X;
X.insert(1);
X.insert(2);
X.insert(4);
X.insert(8);
X.insert(16);

cout<<*X.find_by_order(1)<<endl; // 2
cout<<*X.find_by_order(2)<<endl; // 4
cout<<*X.find_by_order(4)<<endl; // 16
cout<<(end(X)==X.find_by_order(6))<<endl; // true

cout<<X.order_of_key(-5)<<endl;  // 0
cout<<X.order_of_key(1)<<endl;   // 0
cout<<X.order_of_key(3)<<endl;   // 2
cout<<X.order_of_key(4)<<endl;   // 2
cout<<X.order_of_key(400)<<endl; // 5

PS 사용 예시

• BOJ 3653: https://www.acmicpc.net/problem/3653

• 선언부

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

struct Node {
    int priority; // 높을수록 위쪽에
    int id;
    bool operator<(const Node& t) const {
        if (priority != t.priority) return priority < t.priority;
        return id < t.id;
    }
    bool operator>(const Node& t) const {
        return t < *this;
    }
};
#define ordered_set tree<Node, null_type, less<Node>, rb_tree_tag,tree_order_statistics_node_update>

여기서 사용한 구현체는 rb_tree이다. 그런데, 좀 찾아보니 이 구현체를 splay_tree 등으로도 바꿀 수 있는거 같다. 시간 제한이 빡빡하다면 이런 구현체를 바꿔가면서 테스트해봐도 좋을 듯.

• 풀이 사용

void solve() {
    ordered_set S;
    unordered_map<int, ordered_set::point_iterator> M;

    int N, CASE; scanf("%d %d", &N, &CASE);
    for(int i = 1; i <= N; ++i) {
        auto res = S.insert({N - i + 1, i});
        auto it = res.first;
        M.emplace(i, it);
    }

    // for(auto& it: S) {
    //     printf("(%d, %d)\n", it.priority, it.id);
    // }

    int priority = N + 1;

    for(int i = 0; i < CASE; ++i) {
        int c; scanf("%d", &c);
        auto mit = M.find(c);
        auto sit = mit->second;
        M.erase(mit);
        // 아래 부분이 몇 번째 원소인지 빠르게 찾는 부분이다. 일반 set에는 없는 method
        int cdist = S.order_of_key(*sit); //distance(sit, S.end());
        int cid = sit->id;
        S.erase(sit);
        auto res = S.insert({priority, cid});
        auto it = res.first;
        priority++;
        M.emplace(cid, it);
        printf("%d ", N - cdist - 1);
    }
    printf("\n");
}

• BOJ 1572: https://www.acmicpc.net/problem/1572

이 문제의 경우 딱 읽기만 해도 Ordered Set이 사용되면 편하게 풀 수 있겠다는 생각이 든다. 슬라이딩 윈도우 기법으로 원소를 넣고 빼다가, (K - 1) / 2 번째 원소를 빠르게 찾을 수 있기만 하면 쉽게 구현이 되는 문제이다.

• 선언부

#include <bits/stdc++.h>
using namespace std;

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

struct Node {
    int degree;
    int id;
    bool operator<(const Node& t) const {
        if (degree != t.degree) return degree < t.degree;
        return id < t.id;
    }
    bool operator>(const Node& t) const {
        return t < *this;
    }
};
#define ordered_set tree<Node, null_type, less<Node>, rb_tree_tag,tree_order_statistics_node_update>

ordered_set v;

id의 경우 중복값 처리를 위해 존재한다. (중복되서 사라지는 것을 방지)

• 로직 작성

int A[250000];
int main() {
    int N, K; scanf("%d %d", &N, &K);
    for(int i = 0; i < N; ++i) scanf("%d", &A[i]);
    int cur = 0;
    for(int i = 0; i < K; ++i) {
        v.insert({A[i], i});
    }
    int mid = (K - 1) / 2;
    long long ans = v.find_by_order(mid)->degree;

    int l = 0;
    int r = K;
    for(; r < N; ++r, ++l) {
        v.insert({A[r], r});
        v.erase(v.find({A[l], l}));
        ans += v.find_by_order(mid)->degree;
    }
    printf("%lld\n", ans);
    return 0;
}